import matplotlib.pyplot as plt
from sklearn.datasets import fetch_olivetti_faces
from sklearn import decomposition 

dataset = fetch_olivetti_faces(random_state=14)
faces = dataset.data
print(faces.shape)
n_faces, n_hidden_factors = 4, 7
def plot_gallery(images, n_col=n_hidden_factors+1, n_row=n_faces,
                 cmap=plt.cm.gray):
    fig, axs = plt.subplots(n_row,n_col, figsize=(18,7),
                            subplot_kw={'xticks':(), 'yticks':()})
    for i, image in zip(range(n_faces*(n_hidden_factors+1)), images):
        r = int(i/n_col); c = i%n_col
        axs[r,c].imshow(image.reshape((64, 64)), cmap=cmap)
        if (r == 0 and c!=0):
            axs[r,c].set_title('#Factors = ' + str(c*5)) 

images = []
for face in [399,199,60,0]:  # some faces from the 400 dataset samples
    images.append(faces[face])  # draw the original image
    for compressed in range(5,5+n_hidden_factors*5,5):
        estimator = decomposition.NMF(n_components=compressed,
                                      init='nndsvda', tol=5e-3)
        W = estimator.fit_transform(faces)
        # draw the compressed image
        images.append(estimator.inverse_transform(W[face]))    
plot_gallery(images)

plt.show()